from typing import Optional

from fastapi import UploadFile, File, APIRouter, Form

from components.router import CenterRoute
from exceptions import InvalidParam
from ioservice.fileserver import save_file
from models.fileserver import FileEntity, FileUploadConfig

router = APIRouter(route_class=CenterRoute)


@router.post("/fileserver/upload", response_model=FileEntity, response_model_exclude_none=True)
async def update_file(file: UploadFile = File(...),
                      creator: Optional[str] = Form(...),
                      display_name: Optional[str] = Form(None)):
    file_type = FileUploadConfig.get_file_type_by_content_type(file.content_type)
    if file_type is None:
        raise InvalidParam("file", f"unsupported file type: {file.content_type}")
    file_size = sum(read_in_chunks(file.file))
    if FileUploadConfig.size_limit[file_type] < file_size:  # !!!危险操作，内存可能会爆
        raise InvalidParam("file", f"file size [{file_size}] is too large")
    # TODO: validate file format
    file.file.seek(0, 0)
    result = await save_file(file.file, file.filename, file_type, file.content_type, creator, display_name)
    result.fill_url()
    return result


def read_in_chunks(file_object, chunk_size=1024 * 8):
    """Lazy function (generator) to read a file piece by piece.
    Default chunk size: 8K."""
    while True:
        data = file_object.read(chunk_size)
        if not data:
            break
        yield len(data)
